package mq

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/ammyhaber/seckill/pkg/limiter"
)

// PeakShaver 流量削峰器
type PeakShaver struct {
	mq            *RabbitMQ                     // 消息队列客户端
	handlers      map[string]func([]byte) error // 消息处理器映射
	handlersMutex sync.RWMutex                  // 处理器映射的锁
	workerCount   int                           // 工作协程数量
	workerPool    chan struct{}                 // 工作协程池
	limiter       limiter.RateLimiter           // 限流器
	ctx           context.Context               // 上下文
	cancel        context.CancelFunc            // 取消函数
}

// NewPeakShaver 创建流量削峰器
func NewPeakShaver(mq *RabbitMQ, workerCount int, rate float64) *PeakShaver {
	ctx, cancel := context.WithCancel(context.Background())

	return &PeakShaver{
		mq:          mq,
		handlers:    make(map[string]func([]byte) error),
		workerCount: workerCount,
		workerPool:  make(chan struct{}, workerCount),
		limiter:     limiter.NewTokenBucket(rate, workerCount*2), // 令牌桶限流器
		ctx:         ctx,
		cancel:      cancel,
	}
}

// RegisterHandler 注册消息处理器
func (ps *PeakShaver) RegisterHandler(messageType string, handler func([]byte) error) {
	ps.handlersMutex.Lock()
	defer ps.handlersMutex.Unlock()

	ps.handlers[messageType] = handler
}

// Start 启动流量削峰服务
func (ps *PeakShaver) Start() error {
	// 初始化消息队列
	if err := ps.mq.InitOrderQueue(); err != nil {
		return fmt.Errorf("初始化队列失败: %w", err)
	}

	// 启动消费者
	err := ps.mq.ConsumeMessage(ps.ctx, func(data []byte) error {
		// 使用限流器控制消费速率
		if !ps.limiter.Allow() {
			// 如果超过限流阈值，重新入队
			return fmt.Errorf("限流生效，消息将重新入队")
		}

		// 获取工作协程槽位
		select {
		case ps.workerPool <- struct{}{}:
			// 获取到槽位，异步处理消息
			go func() {
				defer func() {
					// 释放工作协程槽位
					<-ps.workerPool
				}()

				// 解析消息类型
				var msg struct {
					Type string          `json:"type"`
					Data json.RawMessage `json:"data"`
				}

				if err := json.Unmarshal(data, &msg); err != nil {
					log.Printf("解析消息失败: %v", err)
					return
				}

				// 查找对应的处理器
				ps.handlersMutex.RLock()
				handler, ok := ps.handlers[msg.Type]
				ps.handlersMutex.RUnlock()

				if !ok {
					log.Printf("未找到消息类型[%s]的处理器", msg.Type)
					return
				}

				// 处理消息
				if err := handler([]byte(msg.Data)); err != nil {
					log.Printf("处理消息[%s]失败: %v", msg.Type, err)
				}
			}()
			return nil
		default:
			// 工作协程池已满，消息重新入队
			return fmt.Errorf("工作协程池已满，消息将重新入队")
		}
	})

	if err != nil {
		return fmt.Errorf("启动消费者失败: %w", err)
	}

	// 启动死信队列消费者
	err = ps.mq.ConsumeDeadLetterMessage(ps.ctx, func(data []byte) error {
		log.Printf("处理死信消息: %s", string(data))

		// 解析消息类型
		var msg struct {
			Type string          `json:"type"`
			Data json.RawMessage `json:"data"`
		}

		if err := json.Unmarshal(data, &msg); err != nil {
			log.Printf("解析死信消息失败: %v", err)
			return nil // 确认消息，不再重试
		}

		// 记录死信消息，可以进行告警或其他处理
		log.Printf("死信消息类型: %s, 数据: %s", msg.Type, string(msg.Data))
		return nil
	})

	if err != nil {
		return fmt.Errorf("启动死信队列消费者失败: %w", err)
	}

	// 启动自动重连
	ps.mq.ReconnectOnFailure()

	return nil
}

// Stop 停止流量削峰服务
func (ps *PeakShaver) Stop() {
	ps.cancel()
	ps.mq.Close()
}

// PublishTask 发布任务到消息队列
func (ps *PeakShaver) PublishTask(ctx context.Context, messageType string, data interface{}) error {
	// 构造消息
	msg := struct {
		Type string      `json:"type"`
		Data interface{} `json:"data"`
		Time int64       `json:"time"`
	}{
		Type: messageType,
		Data: data,
		Time: time.Now().Unix(),
	}

	// 序列化消息
	body, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("序列化消息失败: %w", err)
	}

	// 发布消息
	return ps.mq.PublishMessage(ctx, body)
}
